# -*- coding: utf-8 -*-

# File Name： resize_img
# Description :
# Author : lirui
# create_date： 2022/6/6
# Change Activity:
import cv2
import numpy as np
import torch
import torchvision.transforms.functional as tf
from torch.nn.functional import interpolate

from ml.dl.misc.transform import ndarray_img_to_pil, tensor_to_pil, pil_to_ndarray


def _resize(img, shape):
    """

    Args:
        img:
        shape:

    Returns:

    """
    dtype = img.dtype
    # if dtype != np.uint8:
    #     img = img.astype(np.uint8)
    return cv2.resize(img, shape).astype(dtype)


def _get_shape(img, scale):
    """
    get shape
    Args:
        img:
        scale:

    Returns:

    """
    h, w = img.shape[:2]
    h, w = int(h * scale), int(w * scale)
    return w, h


def resize_img(img, shape=None, scale=None, target=None):
    """
    resize img

    Args:
        target:
        img:
        shape:(h,w)
        scale:

    Returns:

    """
    use_numpy = False
    if isinstance(img, np.ndarray):
        use_numpy = True
        img = ndarray_img_to_pil(img)
    assert shape is not None or scale is not None, "one of `shape` and `scale` cant be None."
    if shape is None:
        shape = _get_shape(img, scale)
    size = shape
    rescaled_image = tf.resize(img, size)
    if target is None:
        if use_numpy:
            return pil_to_ndarray(rescaled_image), None
        return rescaled_image, None
    rescaled_image = tensor_to_pil(rescaled_image)
    ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, img.size))
    ratio_width, ratio_height = ratios
    target = target.copy()
    if "boxes" in target:
        boxes = target["boxes"]
        scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])
        target["boxes"] = scaled_boxes

    if "area" in target:
        area = target["area"]
        scaled_area = area * (ratio_width * ratio_height)
        target["area"] = scaled_area
    h, w = size
    target["size"] = torch.tensor([h, w])

    if "masks" in target:
        target['masks'] = interpolate(
            target['masks'][:, None].float(), size, mode="nearest")[:, 0] > 0.5
    return rescaled_image, target
